/**
 * 
 */
package edu.umd.clip.lm.util.algo;

import java.util.*;
import edu.umd.clip.lm.util.tree.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class BrownAlgorithm <T extends Object> {

	private AnyaryTree<T> theTree;
	private HashMap<AnyaryTree<T>, Integer> elements;
	private HashMap<T, Integer> payloads;
	private Vector<AnyaryTree<T>> elementsByIndex;
	/**
	 *  For the meaning of the variables below, see Peter Brown's paper
	 */
	private double[][] p;
	private double[][] q;
	private double[] pl;
	private double[] pr;
	private double[] s;
	
	public BrownAlgorithm(AnyaryTree<T> tree) {
		theTree = tree;
		elements = new HashMap<AnyaryTree<T>, Integer>(theTree.getChildren().size() * 2);
		payloads = new HashMap<T, Integer>(elements.size());
		
		AnyaryTreeFinder<T> finder = new AnyaryTreeFinder<T>() {
			public boolean apply(AnyaryTree<T> node) {
				if (node.getPayload() != null) {
					Integer num = new Integer(elements.size());
					elements.put(node, num);
					payloads.put(node.getPayload(), num);
				}
				return false;
			}
		};
		
		// populate elements
		theTree.searchPostorder(finder);
		int size = elements.size();
		p = new double[size][size];
		
		elementsByIndex = new Vector<AnyaryTree<T>>(size);
		elementsByIndex.setSize(size);
		for(Map.Entry<AnyaryTree<T>, Integer> pair : elements.entrySet()) {
			elementsByIndex.set(pair.getValue().intValue(), pair.getKey());
		}
	}

	public void setDistribution(T var1, AbstractMap<T, Double> var2) {
		int var1num = payloads.get(var1).intValue();
		double[] arr = p[var1num];
		//pl[var1num] = 0.0;
		for(Map.Entry<T, Double> entry : var2.entrySet()) {
			int var2num = payloads.get(entry.getKey()).intValue();
			double val = entry.getValue().doubleValue();
			arr[var2num] = val;
			//pr[var2num] += val;
			//pl[var1num] += val;
		}
	}

	public AnyaryTree<T> process() {
		int size = elements.size();
		q = new double[size][size];
		s = new double[size];
		pl = new double[size];
		pr = new double[size];
		
		initializeArrays();
		processChildren(theTree);
		
		p = null;
		q = null;
		pl = null;
		pr = null;
		s = null;
		return theTree;
	}
	
	private void initializeArrays() {
		int size = elements.size();
		
		for(int i=0; i<size; ++i) {
			for(int j=0; j<size; ++j) {
				pl[i] += p[i][j];
				pr[j] += p[i][j];
			}
		}

		for(int i=0; i<size; ++i) {
			for(int j=0; j<size; ++j) {
				q[i][j] = p[i][j] == 0.0 ? 0.0 : p[i][j] * Math.log(p[i][j] / (pl[i] * pr[j]));
			}
		}
		
		for (int i=0; i<size; ++i) {
			for(int j=0; j<size; ++j) {
				s[i] += q[j][i] + q[i][j];
			}
			s[i] -= q[i][i];
		}
	}
	
	private void sameValue(double v1, double v2) {
		double diff = Math.abs(v1 - v2);
		double avg = Math.abs(v1 + v2) / 2;
		if (diff * 1e3 > avg) {
			System.out.printf("DIFF: %f vs. %f\n", v1, v2);
			System.out.println(Thread.currentThread().getStackTrace()[2].toString());
			//throw new Error("Different values: " + Double.toString(v1) + " vs. " + Double.toString(v2));
		}
	}
	
	@SuppressWarnings("unused")
	private void testArrays(int[] indices, int numElements) {
		double[][] test_q = new double[numElements][numElements];
		double[][] test_p = new double[numElements][numElements];
		double[] test_s = new double[numElements];
		double[] test_pl = new double[numElements];
		double[] test_pr = new double[numElements];
		
		for(int i=0; i<numElements; ++i) {
			int idx_i = indices[i];
			for(int j=0; j<numElements; ++j) {
				int idx_j = indices[j];
				test_p[i][j] = p[idx_i][idx_j];
			}
		}
		
		for(int i=0; i<numElements; ++i) {
			for(int j=0; j<numElements; ++j) {
				test_pl[i] += test_p[i][j];
				test_pr[j] += test_p[i][j];
			}
		}

		for(int i=0; i<numElements; ++i) {
			for(int j=0; j<numElements; ++j) {
				test_q[i][j] = test_p[i][j] == 0.0 ? 0.0 : test_p[i][j] * Math.log(test_p[i][j] / (test_pl[i] * test_pr[j]));
			}
		}
		
		for (int i=0; i<numElements; ++i) {
			for(int j=0; j<numElements; ++j) {
				test_s[i] += test_q[j][i] + test_q[i][j];
			}
			test_s[i] -= test_q[i][i];
		}
		
		// test arrays
		for(int i=0; i<numElements; ++i) {
			int idx_i = indices[i];
			sameValue(test_pl[i], pl[idx_i]);
			sameValue(test_pr[i], pr[idx_i]);
			for(int j=0; j<numElements; ++j) {
				int idx_j = indices[j];
				sameValue(test_q[i][j], q[idx_i][idx_j]);
			}
		}
		for(int i=0; i<numElements; ++i) {
			int idx_i = indices[i];
			sameValue(test_s[i], s[idx_i]);
		}		
	}
	
	public double sum(double[] arr) {
		double s = 0.0;
		for(double v : arr) { s+= v; }
		return s;
	}
	
	private double estimateMerging(int[] indices, int i, int j, int numElements) {
		int idx_i = indices[i];
		int idx_j = indices[j];
		double L = s[idx_i] + s[idx_j] - q[idx_i][idx_j] - q[idx_j][idx_i];
		
		// q(i+j, i+j)
		double p_ij_ij = p[idx_i][idx_i] + p[idx_i][idx_j] + p[idx_j][idx_i] + p[idx_j][idx_j]; 
		if (p_ij_ij != 0.0) {
			L -= p_ij_ij * Math.log(p_ij_ij / ((pl[idx_i] + pl[idx_j]) * (pr[idx_i] + pr[idx_j])));
		}
		
		for(int m = 0; m<numElements; ++m) {
			if (m == i || m == j) continue;
			int idx_m = indices[m];
			// q(m, i+j)
			double p_m_ij = p[idx_m][idx_i] + p[idx_m][idx_j];
			if (p_m_ij != 0.0) {
				L -= p_m_ij * Math.log(p_m_ij / (pl[idx_m] * (pr[idx_i] + pr[idx_j])));
			}
			
			// q(i+j, m)
			double p_ij_m = p[idx_i][idx_m] + p[idx_j][idx_m];
			if (p_ij_m != 0.0) {
				L -= p_ij_m * Math.log(p_ij_m / ((pl[idx_i] + pl[idx_j]) * pr[idx_m]));
			}
		}
		//System.out.printf("estimate(%d,%d,%d) = %f\n", i, j, numElements, L);
		return L;
	}
	
	@SuppressWarnings("unused")
	private void mergeElements(int[] indices, int i, int j, int numElements) {
		int idx_i = indices[i];
		int idx_j = indices[j];
		System.out.println("merging " + Integer.toString(i) + "+" + Integer.toString(j) + " out of " + Integer.toString(numElements));
		// first, update s
		
		// put the i+j element to i
		//s[idx_i] = s[idx_i] + s[idx_j] - q[idx_i][idx_j] - q[idx_j][idx_i];
		
		int k = numElements-1;
		int idx_k = indices[k];
		for(int m = 0; m<numElements; ++m) {
			if (m == i || m == j) continue;
			int idx_m = indices[m];
			s[idx_m] -= q[idx_m][idx_i] + q[idx_i][idx_m] + q[idx_m][idx_j] + q[idx_j][idx_m];
			// q(m, i+j)
			double p_m_ij = p[idx_m][idx_i] + p[idx_m][idx_j];
			if (p_m_ij != 0.0) {
				s[idx_m] += p_m_ij * Math.log(p_m_ij / (pl[idx_m] * (pr[idx_i] + pr[idx_j])));
			}
			// q(i+j, m)
			double p_ij_m = p[idx_i][idx_m] + p[idx_j][idx_m];
			if (p_ij_m != 0.0) {
				s[idx_m] += p_ij_m * Math.log(p_ij_m / ((pl[idx_i] + pl[idx_j]) * pr[idx_m]));
			}
		}
		// move the last element to vacant place j
		if (j < k) {
			s[idx_j] = s[idx_k] - q[idx_k][idx_i] - q[idx_i][idx_k] - q[idx_k][idx_j] - q[idx_j][idx_k];
			// q(k, i+j)
			double p_k_ij = p[idx_k][idx_i] + p[idx_k][idx_j];
			if (p_k_ij != 0.0) {
				s[idx_j] += p_k_ij * Math.log(p_k_ij / (pl[idx_k] * (pr[idx_i] + pr[idx_j])));
			}
			// q(i+j, k)
			double p_ij_k = p[idx_i][idx_k] + p[idx_j][idx_k];
			if (p_ij_k != 0.0) {
				s[idx_j] += p_ij_k * Math.log(p_ij_k / ((pl[idx_i] + pl[idx_j]) * pr[idx_k]));
			}
		}
		
		// update q
		for(int m = 0; m<numElements; ++m) {
			if (m == i || m == j) continue;
			int idx_m = indices[m];
			// q(m, i+j)
			double p_m_ij = p[idx_m][idx_i] + p[idx_m][idx_j];
			if (p_m_ij != 0.0) {
				q[idx_m][idx_i] = p_m_ij * Math.log(p_m_ij / (pl[idx_m] * (pr[idx_i] + pr[idx_j])));
			} else {
				q[idx_m][idx_i] = 0.0;
			}
			// q(i+j, m)
			double p_ij_m = p[idx_i][idx_m] + p[idx_j][idx_m];
			if (p_ij_m != 0.0) {
				q[idx_i][idx_m] = p_ij_m * Math.log(p_ij_m / ((pl[idx_i] + pl[idx_j]) * pr[idx_m]));
			} else {
				q[idx_i][idx_m] = 0.0;
			}
		}
		double p_ij_ij = p[idx_i][idx_i] + p[idx_i][idx_j] + p[idx_j][idx_i] + p[idx_j][idx_j]; 
		if (p_ij_ij != 0.0) {
			q[idx_i][idx_i] = p_ij_ij * Math.log(p_ij_ij / ((pl[idx_i] + pl[idx_j]) * (pr[idx_i] + pr[idx_j])));
		} else {
			q[idx_i][idx_i] = 0.0;
		}
		if (j < k) {
			q[idx_j] = q[idx_k];
			for(int m = 0; m<k; ++m) {
				int idx_m = indices[m];
				q[idx_m][idx_j] = q[idx_m][idx_k];
			}
		}
		
		// recompute s[i]
		{
			double s_i = 0.0;
			for(int m=0; m<numElements; ++m) {
				int idx_m = indices[m];
				s_i += q[idx_m][idx_i] + q[idx_i][idx_m];
			}
			s[idx_i] = s_i - q[idx_i][idx_i];
		}

		// now update pl and pr
		pl[idx_i] = pl[idx_i] + pl[idx_j];
		pr[idx_i] = pr[idx_i] + pr[idx_j];
		if (j < k) {
			pl[idx_j] = pl[idx_k];
			pr[idx_j] = pr[idx_k];
		}
		
		// finally, update p
		p[idx_i][idx_i] = p[idx_i][idx_i] + p[idx_i][idx_j] + p[idx_j][idx_i] + p[idx_j][idx_j];
		for(int m = 0; m<numElements; ++m) {
			if (m == i || m == j) continue;
			int idx_m = indices[m];
			p[idx_m][idx_i] = p[idx_m][idx_i] + p[idx_m][idx_j];
			p[idx_i][idx_m] = p[idx_i][idx_m] + p[idx_j][idx_m];
		}
		if (j < k) {
			p[idx_j] = p[idx_k];
			for(int m = 0; m<k; ++m) {
				int idx_m = indices[m];
				p[idx_m][idx_j] = p[idx_m][idx_k];
			}
		}
		
	}
	
	private void mergeElements2(final int[] indices, final int i, final int j, final int numElements, final int totalNumElements) {
		int idx_i = indices[i];
		int idx_j = indices[j];
		int k = numElements-1;
		int idx_k = indices[k];

		p[idx_i][idx_i] = p[idx_i][idx_i] + p[idx_i][idx_j] + p[idx_j][idx_i] + p[idx_j][idx_j];

		for(int m = 0; m<totalNumElements; ++m) {
			if (m == i || m == j) continue;
			int idx_m = indices[m];
			p[idx_m][idx_i] = p[idx_m][idx_i] + p[idx_m][idx_j];
			p[idx_i][idx_m] = p[idx_i][idx_m] + p[idx_j][idx_m];
		}
		if (j < k) {
			p[idx_j] = p[idx_k].clone();
			for(int m = 0; m<totalNumElements; ++m) {
				int idx_m = indices[m];
				p[idx_m][idx_j] = p[idx_m][idx_k];
			}
		}
		
		// naively recompute the arrays
		pr = new double[pr.length];
		for(int m=0; m<totalNumElements; ++m) {
			int idx_m = indices[m];
			pl[idx_m] = 0.0;
			for(int l=0; l<totalNumElements; ++l) {
				int idx_l = indices[l];
				pl[idx_m] += p[idx_m][idx_l];
				pr[idx_l] += p[idx_m][idx_l];
			}
		}

		for(int m=0; m<totalNumElements; ++m) {
			int idx_m = indices[m];
			for(int l=0; l<totalNumElements; ++l) {
				int idx_l = indices[l];
				q[idx_m][idx_l] = p[idx_m][idx_l] == 0.0 ? 0.0 : p[idx_m][idx_l] * Math.log(p[idx_m][idx_l] / (pl[idx_m] * pr[idx_l]));
			}
		}
		
		for (int m=0; m<totalNumElements; ++m) {
			int idx_m = indices[m];
			s[idx_m] = 0.0;
			for(int l=0; l<totalNumElements; ++l) {
				int idx_l = indices[l];
				s[idx_m] += q[idx_l][idx_m] + q[idx_m][idx_l];
			}
			s[idx_m] -= q[idx_m][idx_m];
		}
	}
	
	private void processChildren(AnyaryTree<T> tree) {
		if (tree.getChildren() == null) return;
		for(AnyaryTree<T> child : tree.getChildren()) {
			if (!elements.containsKey(child)) {
				processChildren(child);
			}
		}
		
		// now all children have statistical data
		int[] indices = new int[tree.getChildren().size()];
		int[] allIndices = new int[elements.size()];
		{
			//System.out.printf("Processing children, size=%d\n", indices.length);
			if (indices.length == 1) {
				AnyaryTree<T> child = tree.getChildren().iterator().next(); 
				Integer index = elements.get(child);
				elements.remove(child);
				elements.put(tree, index);
				elementsByIndex.set(index.intValue(), tree);
				return;
			}
			int i = 0 /*indices.length - 1*/;
			for(AnyaryTree<T> child : tree.getChildren()) {
				//System.out.printf("%s=%s\n", Integer.toString(i), child.getPayload() == null ? "<NULL>" : child.getPayload().toString());
				int value = elements.get(child).intValue();
				allIndices[i] = value;
				indices[i] = value;
				++i;
			}
			for(Map.Entry<AnyaryTree<T>,Integer> entry : elements.entrySet()) {
				if (!tree.getChildren().contains(entry.getKey())) {
					allIndices[i++] = entry.getValue().intValue();
				}
			}
			//System.out.println();
		}
	
		// remove all children
		tree.getChildren().clear();
		int totalNumElements = allIndices.length;
		
		for(int numElements = indices.length; numElements > 1; --numElements) {
			//System.out.println(Arrays.toString(Arrays.copyOf(indices,numElements)));
			// find two closest items
			double minimalLoss = Double.MAX_VALUE;
			int bestIndex1 = 0, bestIndex2 = 0;
			for(int i=0; i<numElements; ++i) {
				for(int j=i+1; j<numElements; ++j) {
					double loss = estimateMerging(allIndices, i, j, totalNumElements);
					if (loss < minimalLoss) {
						minimalLoss = loss;
						bestIndex1 = i;
						bestIndex2 = j;
					}
				}
			}
			mergeElements2(allIndices, bestIndex1, bestIndex2, numElements, totalNumElements);
			AnyaryTree<T> elem_i = elementsByIndex.get(indices[bestIndex1]);
			AnyaryTree<T> elem_j = elementsByIndex.get(indices[bestIndex2]);
			if (!(elem_i.getPayload() == elem_j.getPayload())) {
				// both non-nulls
				//System.out.println("Merged "+elem_i.getPayload()+" and "+elem_j.getPayload());
			}
			// at the last iteration, use the existing tree 
			AnyaryTree<T> new_elem = numElements > 2 ? new AnyaryTree<T>(null) : tree;
			new_elem.addChild(elem_i);
			new_elem.addChild(elem_j);
			elements.remove(elem_i);
			elements.remove(elem_j);
			elements.put(new_elem, new Integer(indices[bestIndex1]));
			elementsByIndex.set(indices[bestIndex1], new_elem);
			elementsByIndex.set(indices[bestIndex2], elementsByIndex.get(numElements-1));
			//elementsByIndex.setSize(numElements /* -1 */);

			indices[bestIndex2] = allIndices[bestIndex2] = indices[numElements-1];
			--totalNumElements;
			//testArrays(indices, numElements-1);
		}
		//System.out.printf("processed child %s, index=%d\n", tree.toString(), elements.get(tree).intValue());
	}
}
